import torch.nn as nn
import torch.nn.functional as F
from Loss.LossRegister import DefaultLossRegister as LossRegister
import Utils
from Loss.MSE_loss import MSE_loss
import torch


class RangeRouter_Loss(nn.Module):
    def __init__(self, abs_flag, class_nm_threhold, nm_per_class):
        super().__init__()
        self.losser = nn.NLLLoss()
        self.abs_flag = abs_flag
        self.class_nm_threhold = class_nm_threhold
        self.nm_per_class = nm_per_class
        self.eps = 1e-5

    def label_to_class(self, label):
        label_class = Utils.labels_nm_to_class(label, self.abs_flag, self.class_nm_threhold, self.nm_per_class)
        if "cuda" in str(label.device):
            label_class = label_class.cuda()
        return label_class

    def forward(self, pred, label):
        if type(pred) == tuple:
            pred = pred[1]

        pred = torch.log(pred + self.eps)
        label_class = self.label_to_class(label)
        label_class = label_class.flatten()
        loss = self.losser(pred, label_class)
        return loss


# 这个loss其实不太好把握损失部分的比值,难调
@LossRegister.register("MPFNet_loss")
class MPFNet_loss(nn.Module):
    def __init__(self, abs_flag, class_nm_threhold, nm_per_class, range_router_loss_weight=1e4):
        super().__init__()
        self.abs_flag = abs_flag
        self.class_nm_threhold = class_nm_threhold
        self.nm_per_class = nm_per_class
        self.mse_losser = MSE_loss()
        self.range_router_losser = RangeRouter_Loss(abs_flag, class_nm_threhold, nm_per_class)
        self.range_router_loss_weight = range_router_loss_weight
    def forward(self, pred, label):
        mse_loss = self.mse_losser(pred, label)
        range_router_loss = self.range_router_losser(pred, label)
        return mse_loss + range_router_loss * self.range_router_loss_weight


if __name__ == '__main__':
    print(10 % 4)
    print(10 % -4)
    print(-10 % 4)
    print(-10 % -4)
    import torch
    losser = MSE_loss()
    x = torch.randn((3, 1))
    y = torch.randn((3, 1))
    print(losser(x, y))